import os
import sys
import cv2
import random
import numpy as np
#红色方框是预测的位置，蓝色方框是推理的位置

blue = (255, 0, 0)
red = (0, 0, 255)
detector = cv2.CascadeClassifier('F:/PythonProj/haarcascades/haarcascade_frontalface_default.xml')
if __name__ == '__main__':

    # Initialize hog descriptor for people detection
    winSize = (64, 128)
    blockSize = (16, 16)
    blockStride = (8, 8)
    cellSize = (8, 8)
    nbins = 9
    derivAperture = 1
    winSigma = -1
    histogramNormType = 0
    L2HysThreshold = 0.2
    gammaCorrection = True
    nlevels = 64
    signedGradient = False
    #  Load video
    # filename = DATA_PATH + "videos/boy-walking.mp4"
    cap = cv2.VideoCapture(0)
    # Confirm video is open
    if not cap.isOpened():
        print("Unable to read video")
        sys.exit(1)

    # Variable for storing frames
    frameDisplay = []
    # Initialize Kalman filter.
    # Internal state has 6 elements (x, y, width, vx, vy, vw)
    # Measurement has 3 elements (x, y, width ).
    # Note: Height = 2 x width, so it is not part of the state
    # or measurement.
    KF = cv2.KalmanFilter(8, 4, 0)

    # Transition matrix is of the form
    # [
    #   1, 0, 0, dt, 0,  0,
    #   0, 1, 0, 0,  dt, 0,
    #   0, 0, 1, 0,  0,  dt,
    #   0, 0, 0, 1,  0,  0,
    #   0, 0, 0, 0,  1,  0,
    #   0, 0, 0, dt, 0,  1
    # ]
    # because
    # x = x + vx * dt
    # y = y + vy * dt
    # w = y + vw * dt

    # vx = vx
    # vy = vy
    # vw = vw
    KF.transitionMatrix = cv2.setIdentity(KF.transitionMatrix)

    # Measurement matrix is of the form
    # [
    #  1, 0, 0, 0, 0,  0,
    #  0, 1, 0, 0, 0,  0,
    #  0, 0, 1, 0, 0,  0,
    # ]
    # because we are detecting only x, y and w.
    # These quantities are updated.
    KF.measurementMatrix = cv2.setIdentity(KF.measurementMatrix)

    # Variable to store detected x, y and w
    measurement = np.zeros((4, 1), dtype=np.float32)
    # Variables to store detected object and tracked object
    # objectTracked = np.zeros((4, 1), dtype=np.float32)
    # objectDetected = np.zeros((4, 1), dtype=np.float32)

    # Variables to store results of the predict and update (a.k.a correct step).
    updatedMeasurement = np.zeros((4, 1), dtype=np.float32)
    predictedMeasurement = np.zeros((8, 1), dtype=np.float32)

    # Variable to indicate measurement was updated
    measurementWasUpdated = False

    # Timing variable
    ticks = 0
    preTicks = 0

    # Read frames until object is detected for the first time
    success = True
    # dt for Transition matrix
    dt = 0.0
    # Random number generator for randomly selecting frames for update
    random.seed(42)
    while success:
        sucess, frame = cap.read()
        gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
        faces = detector.detectMultiScale(gray, 1.3, 5)
        if len(faces) > 0:
            for (x, y, w, h) in faces:
                cv2.rectangle(frame, (x, y), (x + w, y + h), (255, 0, 0), 2)
                updatedMeasurement = KF.predict()
                predict_x = int(updatedMeasurement[0][0])
                predict_y = int(updatedMeasurement[1][0])
                predict_w = int(updatedMeasurement[2][0])
                predict_h = int(updatedMeasurement[3][0])
                if predict_x>0 or predict_y>0 or predict_w>0 or predict_h>0:
                    cv2.rectangle(frame, (predict_x, predict_y), (predict_x + predict_w, predict_y + predict_h),
                              (0, 0, 255), 2)
                measurement = np.array([[np.float32(x)], [np.float32(y)], [np.float32(w)], [np.float32(h)]])
                # Perform Kalman update step
                KF.correct(measurement)
                #
                # print("-----------------")
                # print(measurement)
                # print(updatedMeasurement)
                # print("-----------------")

        # Display result.
        cv2.imshow("object tracker", frame)
        # Break if ESC pressed
        if cv2.waitKey(5) & 0xFF == 27:
            break

    cap.release()
    cv2.destroyAllWindows()
